"""请求限流中间件"""

import time
import logging
from typing import Dict, Optional, Union
from fastapi import Request, HTTPException
from starlette.middleware.base import BaseHTTPMiddleware
from starlette.responses import Response
from collections import defaultdict, deque

from backend.common.config import config_manager
from backend.middleware.trace_middleware import get_trace_id
from backend.logger_setup import get_logger_with_trace_id

logger = logging.getLogger('middleware.rate_limit')

class RateLimitMiddleware(BaseHTTPMiddleware):
    """请求限流中间件，限制API请求的频率"""
    
    def __init__(self, app, 
                 global_rate_limit: Optional[Dict[str, int]] = None,
                 path_specific_limits: Optional[Dict[str, Dict[str, int]]] = None):
        """初始化限流中间件
        
        Args:
            app: FastAPI应用实例
            global_rate_limit: 全局限流配置，格式为 {"requests": 100, "seconds": 60}，表示60秒内最多100个请求
            path_specific_limits: 路径特定的限流配置，格式为 {"/api/chat": {"requests": 50, "seconds": 60}}
        """
        super().__init__(app)
        
        # 默认限流配置
        default_global_limit = {"requests": 100, "seconds": 60}
        default_path_limits = {}
        
        # 尝试从配置中获取限流设置，如果没有则使用默认值
        try:
            self.global_rate_limit = global_rate_limit or config_manager.rate_limit.global_limit
            self.path_specific_limits = path_specific_limits or config_manager.rate_limit.path_limits
        except (AttributeError, TypeError):
            # 如果配置不存在或格式不正确，使用默认值
            self.global_rate_limit = global_rate_limit or default_global_limit
            self.path_specific_limits = path_specific_limits or default_path_limits
            logger.warning(f"无法从配置中获取限流设置，使用默认值")
        
        # 存储每个客户端的请求时间
        # 格式: {client_key: {path_pattern: deque([timestamp1, timestamp2, ...])}}
        self.request_times = defaultdict(lambda: defaultdict(deque))
        
        # 存储路径模式与限流配置的映射
        self.path_limit_mapping = self._build_path_limit_mapping()
    
    def _build_path_limit_mapping(self) -> Dict[str, Dict[str, int]]:
        """构建路径模式与限流配置的映射
        
        Returns:
            路径模式与限流配置的映射
        """
        path_mapping = {}
        
        # 添加路径特定的限流配置
        for path_pattern, limits in self.path_specific_limits.items():
            path_mapping[path_pattern] = limits
        
        return path_mapping
    
    def _get_client_key(self, request: Request) -> str:
        """获取客户端的唯一标识
        
        Args:
            request: FastAPI请求对象
            
        Returns:
            客户端唯一标识
        """
        # 优先使用X-Forwarded-For头（如果有）
        x_forwarded_for = request.headers.get('X-Forwarded-For')
        if x_forwarded_for:
            # 取第一个IP地址
            return x_forwarded_for.split(',')[0].strip()
        
        # 其次使用客户端IP地址
        client_host = request.client.host
        if client_host:
            return client_host
        
        # 最后使用默认值
        return 'unknown_client'
    
    def _get_path_pattern(self, path: str) -> str:
        """获取路径对应的模式
        
        Args:
            path: 请求路径
            
        Returns:
            匹配的路径模式
        """
        # 首先尝试精确匹配
        if path in self.path_limit_mapping:
            return path
        
        # 然后尝试前缀匹配
        for pattern in sorted(self.path_limit_mapping.keys(), key=len, reverse=True):
            if path.startswith(pattern) and (pattern.endswith('/') or path[len(pattern):].startswith('/')):
                return pattern
        
        # 如果没有匹配的路径模式，返回全局模式
        return 'global'
    
    def _is_rate_limited(self, client_key: str, path: str) -> bool:
        """检查客户端是否被限流
        
        Args:
            client_key: 客户端唯一标识
            path: 请求路径
            
        Returns:
            是否被限流
        """
        current_time = time.time()
        path_pattern = self._get_path_pattern(path)
        
        # 获取限流配置，如果没有特定路径的配置，则使用全局配置
        limits = self.path_limit_mapping.get(path_pattern, self.global_rate_limit)
        max_requests = limits.get('requests', 100)
        time_window = limits.get('seconds', 60)
        
        # 获取或创建请求时间队列
        request_queue = self.request_times[client_key][path_pattern]
        
        # 移除过期的请求时间
        while request_queue and current_time - request_queue[0] > time_window:
            request_queue.popleft()
        
        # 检查是否超过限制
        if len(request_queue) >= max_requests:
            return True
        
        # 添加当前请求时间
        request_queue.append(current_time)
        return False
    
    async def dispatch(self, request: Request, call_next) -> Response:
        # 获取trace_id
        trace_id = get_trace_id(request)
        
        # 获取带trace_id的logger
        log = get_logger_with_trace_id(logger, trace_id)
        
        # 跳过不需要限流的路径
        if request.url.path.startswith('/static/'):
            return await call_next(request)
        
        # 获取客户端标识和路径
        client_key = self._get_client_key(request)
        path = request.url.path
        
        # 检查是否被限流
        if self._is_rate_limited(client_key, path):
            path_pattern = self._get_path_pattern(path)
            limits = self.path_limit_mapping.get(path_pattern, self.global_rate_limit)
            
            log.warning(f"请求被限流 [trace_id={trace_id}, client={client_key}, path={path}]")
            
            # 返回429 Too Many Requests
            return Response(
                content=f"Too many requests. Please try again later.",
                status_code=429,
                headers={
                    "X-Trace-Id": trace_id,
                    "Retry-After": str(limits.get('seconds', 60))
                }
            )
        
        # 执行请求处理
        response = await call_next(request)
        
        # 添加限流相关的响应头
        path_pattern = self._get_path_pattern(path)
        request_queue = self.request_times[client_key][path_pattern]
        limits = self.path_limit_mapping.get(path_pattern, self.global_rate_limit)
        
        remaining = max(0, limits.get('requests', 100) - len(request_queue))
        response.headers['X-RateLimit-Limit'] = str(limits.get('requests', 100))
        response.headers['X-RateLimit-Remaining'] = str(remaining)
        response.headers['X-RateLimit-Reset'] = str(int(time.time() + limits.get('seconds', 60) - (time.time() - request_queue[0] if request_queue else 0)))
        
        return response

def setup_rate_limiting(app):
    """设置请求限流
    
    Args:
        app: FastAPI应用实例
    """
    # 默认限流配置
    default_global_limit = {"requests": 100, "seconds": 60}
    default_path_limits = {}
    
    # 尝试从配置中获取限流设置，如果没有则使用默认值
    try:
        rate_limit_config = config_manager.rate_limit
        global_limit = rate_limit_config.global_limit
        path_limits = rate_limit_config.path_limits
    except (AttributeError, TypeError):
        # 如果配置不存在或格式不正确，使用默认值
        global_limit = default_global_limit
        path_limits = default_path_limits
        logger.warning(f"无法从配置中获取限流设置，使用默认值")
    
    # 注册限流中间件
    app.add_middleware(
        RateLimitMiddleware,
        global_rate_limit=global_limit,
        path_specific_limits=path_limits
    )

__all__ = [
    'RateLimitMiddleware',
    'setup_rate_limiting'
]